#define WIN32_LEAN_AND_MEAN
#include <Windows.h>
#include <DbgHelp.h>

#include "StackWalker.h"
#include "UnwindStructures.h"

#include <stdio.h>
#include <assert.h>

#ifdef OLD_WINDOWS_SDK_VERSION

PEXCEPTION_ROUTINE (WINAPI *RtlVirtualUnwind)( ULONG ,
											   ULONG64 ,
											   ULONG64 ,
											   PRUNTIME_FUNCTION ,
											   PCONTEXT ,
											   PVOID *,
											   PULONG64 ,
											   PKNONVOLATILE_CONTEXT_POINTERS ) = 0;

PRUNTIME_FUNCTION (WINAPI *RtlLookupFunctionEntry)(	ULONG64,
													PULONG64,
													PUNWIND_HISTORY_TABLE HistoryTable ) = 0;

#endif


int WalkStack(int zStackDepth, StackTrace& zStackTrace)
{
	CONTEXT                       context;
	KNONVOLATILE_CONTEXT_POINTERS NvContext;
	PRUNTIME_FUNCTION             runtimeFunction;
	PVOID                         handlerData;
	ULONG64                       establisherFrame;
	ULONG64                       imageBase;
	IMAGEHLP_SYMBOL				  sym;
	HANDLE						  currentProcess;

	currentProcess = GetCurrentProcess();

	// Initialiazation
	if(RtlVirtualUnwind == NULL && RtlLookupFunctionEntry == NULL)
	{
		const HMODULE hNtDll = ::GetModuleHandle("ntdll.dll");
		reinterpret_cast<void*&>(RtlVirtualUnwind) = ::GetProcAddress(hNtDll, "RtlVirtualUnwind");
		reinterpret_cast<void*&>(RtlLookupFunctionEntry) = ::GetProcAddress(hNtDll, "RtlLookupFunctionEntry");

		// Load symbols.
		if(!SymInitialize (currentProcess, NULL, TRUE))
        {
            // typically this happens because something else called SymInitialize
            // if this occurs then at some point calling SymGetSymFromAddr will
            // cause a stack overwrite. Ugh!
            const unsigned int error = GetLastError();
            printf("SymInitialize failed with error %u\n", error);
            assert(!"SymInitialize failed");
        }
	}

	sym.SizeOfStruct  = sizeof (sym);
	sym.MaxNameLength = 0;	

	//First, we'll get the caller's context.
	RtlCaptureContext(&context);

	// This unwind loop intentionally skips the first call frame, as it shall
	// correspond to the call to StackTrace64, which we aren't interested in.
	int frame = 0;
	
	while(true)
	{
		// Try to look up unwind metadata for the current function.
		runtimeFunction = RtlLookupFunctionEntry(context.Rip, &imageBase, NULL);

		// The first frame is the current function so we skip it. 
		if(frame > 0)
		{
            assert((frame - 1) >= 0);
            assert((frame - 1) < MAX_STACK_TRACE_SIZE);
                
			zStackTrace[frame-1].rsp = context.Rsp;
			zStackTrace[frame-1].rip = context.Rip;
		}

		if (!runtimeFunction)
		{
			//If we don't have a RUNTIME_FUNCTION, then we've encountered
			//a leaf function.  Adjust the stack appropriately.
			context.Rip  = (ULONG64)(*(PULONG64)context.Rsp);
			context.Rsp += 8;
		}
		else
		{
			//Otherwise, call upon RtlVirtualUnwind to execute the unwind for us.
			RtlZeroMemory(&NvContext, sizeof(KNONVOLATILE_CONTEXT_POINTERS));
			RtlVirtualUnwind(0,	imageBase, context.Rip, runtimeFunction, &context, &handlerData,
				&establisherFrame,	&NvContext);
		}

		//If we reach an RIP of zero, this means that we've walked off the end
		//of the call stack and are done.
		if (!context.Rip)
			break;

		// The first frame is the current function so we skip it. 
		if(frame == zStackDepth)
		{
            assert((frame - 1) >= 0);
            assert((frame - 1) < MAX_STACK_TRACE_SIZE);

			zStackTrace[frame-1].rspBeforePrologue = context.Rsp - 8;
			zStackTrace[frame-1].returnAddress     = context.Rip;

            assert(sym.MaxNameLength == 0);			
		
			// SymGetSymFromAddr can be quite slow so we only want to call it when there is no other way
			// to get a function address.
			if(runtimeFunction)
			{
				// Non-leaf functions have unwind data so we can retrieve their address from the information
				// contained in PRUNTIME_FUNCTION.
				zStackTrace[frame-1].functionAddress = imageBase + runtimeFunction->BeginAddress;
			}
			else
			{
				// The address of leaf functions is retrieved from the debug symbols.
				if (SymGetSymFromAddr(currentProcess, zStackTrace[frame-1].rip, NULL, &sym)) 
				{
					zStackTrace[frame-1].functionAddress = sym.Address;
				}
			}			
			
			return frame;	
		}		

		frame++;
	}

	return frame;
}



